package com.lkdont.android.media

import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.media.MediaMuxer
import android.util.Log
import android.view.Surface
import java.nio.ByteBuffer

/**
 * 多媒体编码器
 *
 * @author lkdont
 */
class MediaEncoder {

    companion object {
        private const val TAG = "MediaEncoder"
    }

    /*** 编码器参数 */
    class Config private constructor(
        val outputPath: String,
        val outputFormat: Int,
        // 音频参数
        val audioFormat: String?,
        val audioChannelCount: Int,
        val audioSampleRate: Int,
        val audioBitrate: Int,
        // 视频参数
        val videoFormat: String?,
        val videoWidth: Int,
        val videoHeight: Int,
        val videoBitrate: Int,
        val videoFrameRate: Int,
        val videoIFrameInterval: Int
    ) {

        class Builder(val outputPath: String, val outputFormat: Int) {
            // 音频参数
            var audioFormat: String? = null
            fun audioFormat(audioFormat: String?) = apply { this.audioFormat = audioFormat }

            var audioChannelCount: Int = 0
            fun audioChannelCount(audioChannelCount: Int) =
                apply { this.audioChannelCount = audioChannelCount }

            var audioSampleRate: Int = 0
            fun audioSampleRate(audioSampleRate: Int) =
                apply { this.audioSampleRate = audioSampleRate }

            var audioBitrate: Int = 0
            fun audioBitrate(audioBitrate: Int) = apply { this.audioBitrate = audioBitrate }

            // 视频参数
            var videoFormat: String? = null
            fun videoFormat(videoFormat: String?) = apply { this.videoFormat = videoFormat }

            var videoWidth: Int = 0
            fun videoWidth(videoWidth: Int) = apply { this.videoWidth = videoWidth }

            var videoHeight: Int = 0
            fun videoHeight(videoHeight: Int) = apply { this.videoHeight = videoHeight }

            var videoBitrate: Int = 0
            fun videoBitrate(videoBitrate: Int) = apply { this.videoBitrate = videoBitrate }

            var videoFrameRate: Int = 0
            fun videoFrameRate(videoFrameRate: Int) = apply { this.videoFrameRate = videoFrameRate }

            var videoIFrameInterval: Int = 0
            fun videoIFrameInterval(videoIFrameInterval: Int) =
                apply { this.videoIFrameInterval = videoIFrameInterval }

            fun build() = Config(
                outputPath,
                outputFormat,
                audioFormat,
                audioChannelCount,
                audioSampleRate,
                audioBitrate,
                videoFormat,
                videoWidth,
                videoHeight,
                videoBitrate,
                videoFrameRate,
                videoIFrameInterval
            )
        }

    }

    private var mediaMuxer: MediaMuxer? = null

    // 音频
    private var audioCodec: MediaCodec? = null
    private var audioBufferInfo: MediaCodec.BufferInfo? = null
    private var audioTrackId: Int = -1

    // 视频
    private var videoCodec: MediaCodec? = null
    private var videoBufferInfo: MediaCodec.BufferInfo? = null
    private var videoTrackId: Int = -1

    /*** 视频输入Surface */
    var videoInputSurface: Surface? = null
        private set

    private var isAudioEnable = false
    private var isVideoEnable = false

    val isMediaMuxerReady: Boolean
        get() {
            return ((!isAudioEnable || audioTrackId >= 0) && (!isVideoEnable || videoTrackId >= 0))
        }


    private lateinit var encoderConfig: Config

    fun start(config: Config) {
        encoderConfig = config
        val audioFormat = config.audioFormat
        val videoFormat = config.videoFormat
        isAudioEnable = (audioFormat != null)
        isVideoEnable = (videoFormat != null)
        if (!isAudioEnable && !isVideoEnable) {
            throw MediaException("There is nothing to encode.")
        }
        // 初始化音频编码器
        if (isAudioEnable) {
            if (audioBufferInfo == null) {
                audioBufferInfo = MediaCodec.BufferInfo()
            }
            audioCodec = MediaCodec.createEncoderByType(audioFormat!!)
            val format = MediaFormat.createAudioFormat(
                audioFormat,
                config.audioSampleRate,
                config.audioChannelCount
            )
            format.setInteger(MediaFormat.KEY_BIT_RATE, config.audioBitrate)
            audioCodec!!.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
            audioCodec!!.start()
        }
        // 初始化视频编码器
        if (isVideoEnable) {
            if (videoBufferInfo == null) {
                videoBufferInfo = MediaCodec.BufferInfo()
            }
            videoCodec = MediaCodec.createEncoderByType(videoFormat!!)
            val format =
                MediaFormat.createVideoFormat(videoFormat, config.videoWidth, config.videoHeight)
            format.setInteger(
                MediaFormat.KEY_COLOR_FORMAT,
                MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface
            )
            format.setInteger(MediaFormat.KEY_BIT_RATE, config.videoBitrate)
            format.setInteger(MediaFormat.KEY_FRAME_RATE, config.videoFrameRate)
            format.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, config.videoIFrameInterval)
            videoCodec!!.configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
            videoInputSurface = videoCodec!!.createInputSurface()
            videoCodec!!.start()
        }
        // 初始化包装器
        mediaMuxer = MediaMuxer(config.outputPath, config.outputFormat)
        audioTrackId = -1
        videoTrackId = -1
    }

    /*** 输入音频数据 */
    fun writeAudioData(presentationTimeUs: Long, data: ByteArray, length: Int) {
        val codec = audioCodec ?: return
        writeData(codec, presentationTimeUs, data, length, false)
    }

    /*** 输入数据 */
    private fun writeData(
        codec: MediaCodec,
        presentationTimeUs: Long,
        data: ByteArray,
        dataLength: Int,
        endOfStream: Boolean
    ) {
        // 获取编码器的输入流
        val inputId: Int = codec.dequeueInputBuffer(50)
        if (inputId >= 0) {
            val inputBuffer: ByteBuffer? = codec.getInputBuffer(inputId)
            if (inputBuffer != null) {
                // 将数据放入编码队列
                inputBuffer.put(data, 0, dataLength)
                codec.queueInputBuffer(
                    inputId,
                    0,
                    dataLength,
                    presentationTimeUs,
                    if (endOfStream) MediaCodec.BUFFER_FLAG_END_OF_STREAM else 0
                )
            }
        } else {
            Log.w(TAG, "writeData: 获取编码器的输入流失败")
        }
    }

    private fun drainEncoder(
        encoder: MediaCodec,
        bufferInfo: MediaCodec.BufferInfo,
        endOfStream: Boolean
    ) {
        if (endOfStream) {
            if (encoder == videoCodec && videoInputSurface != null) {
                // 当视频编码器使用Surface作为输入源的时候才调用此方法
                encoder.signalEndOfInputStream()
            }
//            else {
//                // 写入一个空的数据帧，提示数据输出已经结束
//                if (encoder == audioCodec) {
//                    writeData(encoder, 0, ByteArray(0), true)
//                }
//            }
        }
        while (true) {
            val status = encoder.dequeueOutputBuffer(bufferInfo, 0)
            when {
                status == MediaCodec.INFO_TRY_AGAIN_LATER -> {
                    // 现在没有数据要输出
//                    if (!endOfStream) {
                    return
//                    } else {
//                        Log.d(
//                            TAG,
//                            "no output available, spinning to await EOS"
//                        )
//                    }
                }
                status == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED -> {
                    // 设置MediaMuxer
                    val format = encoder.outputFormat
                    Log.d(
                        TAG,
                        "encoder output format changed: $format"
                    )
                    if (encoder == audioCodec && audioTrackId < 0) {
                        audioTrackId = mediaMuxer!!.addTrack(format)
                    } else if (encoder == videoCodec && videoTrackId < 0) {
                        videoTrackId = mediaMuxer!!.addTrack(format)
                    } else {
                        throw MediaException("MediaMuxer addTrack fail.")
                    }
                    if (isMediaMuxerReady) {
                        mediaMuxer!!.start()
                    }
                }
                status < 0 -> {
                    // do nothing
                    Log.w(
                        TAG,
                        "unexpected result from encoder.dequeueOutputBuffer: $status"
                    )
                }
                else -> {
                    if (!isMediaMuxerReady) {
                        Log.w(TAG, "drainEncoder: muxer hasn't started")
                        return
                    }
                    val trackIndex = if (encoder == audioCodec) audioTrackId else videoTrackId;
                    val encodedData =
                        encoder.getOutputBuffer(status) ?: throw MediaException(
                            "Encode fail."
                        )
                    if (bufferInfo.flags and MediaCodec.BUFFER_FLAG_CODEC_CONFIG != 0) {
                        // The codec config data was pulled out and fed to the muxer when we got
                        // the INFO_OUTPUT_FORMAT_CHANGED status.  Ignore it.
                        Log.d(
                            TAG,
                            "ignoring BUFFER_FLAG_CODEC_CONFIG"
                        )
                        bufferInfo.size = 0
                    }
                    if (bufferInfo.size != 0) {
                        // adjust the ByteBuffer values to match BufferInfo (not needed?)
                        encodedData.position(bufferInfo.offset)
                        encodedData.limit(bufferInfo.offset + bufferInfo.size)
                        mediaMuxer!!.writeSampleData(trackIndex, encodedData, bufferInfo)
//                        Log.d(
//                            TAG,
//                            "($status)sent " + bufferInfo.size + " bytes to muxer, ts=" +
//                                    bufferInfo.presentationTimeUs
//                        )
                    }

                    encoder.releaseOutputBuffer(status, false)

                    if (bufferInfo.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM != 0) {
                        if (!endOfStream) {
                            Log.w(
                                TAG,
                                "reached end of stream unexpectedly"
                            )
                        } else {
                            Log.d(
                                TAG,
                                "end of stream reached"
                            )
                        }
                        return // out of while
                    }
                }
            }
        }
    }

    fun drainAudioEncoder(endOfStream: Boolean) {
        if (isAudioEnable) {
            drainEncoder(audioCodec!!, audioBufferInfo!!, endOfStream)
        }
    }

    fun drainVideoEncoder(endOfStream: Boolean) {
        if (isVideoEnable) {
            drainEncoder(videoCodec!!, videoBufferInfo!!, endOfStream)
        }
    }

    fun stop() {
        if (audioCodec != null) {
            audioCodec!!.stop()
            audioCodec!!.release()
            audioCodec = null
        }
        if (videoCodec != null) {
            videoCodec!!.stop()
            videoCodec!!.release()
            videoCodec = null
        }
        if (mediaMuxer != null) {
            mediaMuxer!!.stop()
            mediaMuxer!!.release()
            mediaMuxer = null
        }
        // 重置变量
        audioTrackId = -1
        isAudioEnable = false
        videoTrackId = -1
        isVideoEnable = false
    }
}